#include "ext4_scanner.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <ext2fs/ext2fs.h>
#include <dirent.h> // 引入用于目录遍历的头文件

// Helper: 手动组合 64 位文件大小 (不变)
static inline __u64 get_inode_size(const struct ext2_inode *inode) {
    if (LINUX_S_ISREG(inode->i_mode)) {
        return ((__u64) inode->i_size_high << 32) | inode->i_size;
    }
    return inode->i_size;
}

// Helper: 写入 blob (不变)
static int write_chunk_to_blob(const char* output_dir, const char* hash, const unsigned char* data) {
    char blob_path[1024];
    snprintf(blob_path, sizeof(blob_path), "%s/blobs/sha256/%s", output_dir, hash);
    FILE* fp = fopen(blob_path, "wb");
    if (!fp) { LOG_ERROR("无法打开 blob 文件写入: %s", blob_path); return -1; }
    fwrite(data, 1, CHUNK_SIZE, fp);
    fclose(fp);
    return 0;
}

// Forward declaration
static errcode_t recursive_dir_scan(ext2_filsys fs, ext2_ino_t dir_ino, callback_context_t* ctx);

// 处理单个文件 Inode (不变)
static void process_file_inode(ext2_filsys fs, ext2_ino_t ino, struct ext2_inode *inode_ptr, callback_context_t* ctx) {
    char buf[CHUNK_SIZE];
    blk64_t num_blocks = (get_inode_size(inode_ptr) + CHUNK_SIZE - 1) / CHUNK_SIZE;

    for (blk64_t logical_block = 0; logical_block < num_blocks; logical_block++) {
        blk64_t physical_block = 0;
        if (ext2fs_bmap2(fs, ino, inode_ptr, NULL, 0, logical_block, 0, &physical_block) != 0) continue;
        if (physical_block == 0) continue;

        (*ctx->processed_blocks)++;

        if (io_channel_read_blk64(fs->io, physical_block, 1, buf) != 0) continue;

        char hash_str[XXH64_HEX_LEN + 1];
        calculate_xxh64_hex((unsigned char*)buf, CHUNK_SIZE, hash_str);
        
        if (!g_hash_table_contains(ctx->seen_chunks_set, hash_str)) {
            (*ctx->unique_blocks)++;
            if (write_chunk_to_blob(ctx->output_dir, hash_str, (unsigned char*)buf) == 0) {
                g_hash_table_add(ctx->seen_chunks_set, g_strdup(hash_str));
            }
        }
        
        lba_map_entry_t entry;
        entry.lba = physical_block; 
        strncpy(entry.hash, hash_str, XXH64_HEX_LEN);
        entry.hash[XXH64_HEX_LEN] = '\0';
        fwrite(&entry, sizeof(lba_map_entry_t), 1, ctx->lba_map_file);
    }
}

// ext2fs_dir_iterate 的回调函数 (不变)
static int dir_iterate_callback(struct ext2_dir_entry *dirent,
                                int offset,
                                int blocksize,
                                char *buf,
                                void *priv_data) {
    callback_context_t *ctx = (callback_context_t *)priv_data;
    
    char name[EXT2_NAME_LEN + 1];
    int namelen = dirent->name_len & 0xFF;
    if (namelen > EXT2_NAME_LEN) namelen = EXT2_NAME_LEN;
    strncpy(name, dirent->name, namelen);
    name[namelen] = '\0';

    if (dirent->inode == 0 || strcmp(name, ".") == 0 || strcmp(name, "..") == 0) {
        return 0;
    }

    struct ext2_inode entry_inode;
    if (ext2fs_read_inode(ctx->fs, dirent->inode, &entry_inode) != 0) {
        return 0;
    }
    
    if (LINUX_S_ISREG(entry_inode.i_mode)) {
        process_file_inode(ctx->fs, dirent->inode, &entry_inode, ctx);
    } else if (LINUX_S_ISDIR(entry_inode.i_mode)) {
        recursive_dir_scan(ctx->fs, dirent->inode, ctx);
    }
    
    return 0;
}

// 递归遍历目录的函数 (不变)
static errcode_t recursive_dir_scan(ext2_filsys fs, ext2_ino_t dir_ino, callback_context_t* ctx) {
    errcode_t retval;
    char *block_buf = NULL;

    retval = ext2fs_get_mem(fs->blocksize, &block_buf);
    if (retval) return retval;

    retval = ext2fs_dir_iterate(fs, dir_ino, 0, block_buf, dir_iterate_callback, ctx);

    ext2fs_free_mem(&block_buf);
    return retval;
}

int scan_ext4_and_build_map(const char *device_path, const char* output_dir) {
    LOG_INFO("阶段 2: 扫描文件系统并构建 LBA Map...");
    
    errcode_t retval;
    ext2_filsys fs;

    GHashTable* seen_chunks_set = g_hash_table_new_full(g_str_hash, g_str_equal, g_free, NULL);
    if (!seen_chunks_set) { LOG_ERROR("创建哈希表失败"); return -1; }

    char blob_dir[1024], lba_map_path[1024];
    snprintf(blob_dir, sizeof(blob_dir), "%s/blobs/sha256", output_dir);
    snprintf(lba_map_path, sizeof(lba_map_path), "%s/lba_map.bin", output_dir);
    if (run_command("mkdir -p %s", blob_dir) != 0) { LOG_ERROR("创建 blobs 目录失败"); g_hash_table_destroy(seen_chunks_set); return -1; }
    
    // ***** 关键修改：预加载已有的 blobs *****
    LOG_INFO("预加载已存在的 blobs...");
    DIR *d;
    struct dirent *dir;
    d = opendir(blob_dir);
    if (d) {
        while ((dir = readdir(d)) != NULL) {
            // 跳过 "." 和 ".."
            if (strcmp(dir->d_name, ".") != 0 && strcmp(dir->d_name, "..") != 0) {
                // 将文件名（即哈希）添加到集合中
                g_hash_table_add(seen_chunks_set, g_strdup(dir->d_name));
            }
        }
        closedir(d);
    } else {
        // 目录不存在是正常情况（第一次运行时）
        LOG_INFO("blobs 目录不存在, 将从零开始创建。");
    }
    LOG_INFO("预加载完成, 已加载 %u 个 blobs。", g_hash_table_size(seen_chunks_set));


    LOG_INFO("准备使用 libext2fs 打开设备: %s", device_path);
    retval = ext2fs_open(device_path, 0, 0, 0, unix_io_manager, &fs);
    if (retval) { com_err("ext2fs_open", retval, "无法打开设备 %s", device_path); g_hash_table_destroy(seen_chunks_set); return -1; }
    LOG_INFO("ext2fs_open 成功！");
    
    FILE* lba_map_file = fopen(lba_map_path, "wb");
    if (!lba_map_file) { LOG_ERROR("无法创建 LBA Map 文件"); ext2fs_close(fs); g_hash_table_destroy(seen_chunks_set); return -1; }

    LOG_INFO("开始通过目录递归扫描文件...");
    long long processed_blocks = 0;
    long long unique_blocks = 0; // 这个计数现在代表“本次新增”的唯一块

    callback_context_t context = {
        .fs = fs,
        .output_dir = output_dir, 
        .lba_map_file = lba_map_file, 
        .seen_chunks_set = seen_chunks_set,
        .processed_blocks = &processed_blocks, 
        .unique_blocks = &unique_blocks
    };
    
    recursive_dir_scan(fs, EXT2_ROOT_INO, &context);
    
    printf("\n");

    LOG_INFO("目录扫描完成。共处理 %lld 个数据块, 本次新增 %lld 个唯一块。", processed_blocks, unique_blocks);
    double dedup_ratio = processed_blocks > 0 ? (1.0 - (double)unique_blocks / processed_blocks) * 100.0 : 0;
    LOG_INFO("本次运行的去重率 (与已存在 blobs 比较): %.2f%%", dedup_ratio);

    fclose(lba_map_file);
    ext2fs_close(fs);
    g_hash_table_destroy(seen_chunks_set);

    LOG_INFO("阶段 2 完成: LBA Map 和去重后的 blobs 已生成。");
    return 0;
}